# coding: UTF-8
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn import metrics
import time
from utils import get_time_dif
from tensorboardX import SummaryWriter


# 权重初始化，默认xavier
def init_network(model, method='xavier', exclude='embedding', seed=123):
    for name, w in model.named_parameters():
        if exclude not in name:
            if 'weight' in name:
                if method == 'xavier':
                    nn.init.xavier_normal_(w)
                elif method == 'kaiming':
                    nn.init.kaiming_normal_(w)
                else:
                    nn.init.normal_(w)
            elif 'bias' in name:
                nn.init.constant_(w, 0)
            else:
                pass

def my_cross_entropy(x1, x2):
    x1 = F.softmax(x1, dim=1)
    x2 = F.softmax(x2, dim=1)
    return -torch.sum(x2 * (x1 ** 2) * x1.log(), dim=1)


# Rewritten Loss functions
def my_loss_coteaching(y_1, y_2, t, forget_rate):
    # self small loss and other teaches

    loss_1_2 = my_cross_entropy(y_1, y_2.detach())
    loss_2_1 = my_cross_entropy(y_2, y_1.detach())
    loss_1_t = F.cross_entropy(y_1, t, reduction = 'none')
    loss_2_t = F.cross_entropy(y_2, t, reduction = 'none')
    loss_1_t_sorted, ind_1_t_sorted = torch.sort(loss_1_t)
    loss_2_t_sorted, ind_2_t_sorted = torch.sort(loss_2_t)
    
    remember_rate = 1 - forget_rate
    num_remember = int(remember_rate * len(t))

    # exchange
    loss_1_update = torch.cat((loss_1_t_sorted[:num_remember], loss_1_2[ind_1_t_sorted[num_remember:]]))
    loss_2_update = torch.cat((loss_2_t_sorted[:num_remember], loss_2_1[ind_2_t_sorted[num_remember:]]))

    return torch.mean(loss_1_update), torch.mean(loss_2_update)

# Rewritten Loss functions
def my_loss_coteaching_2(y_1, y_2, t, forget_rate):
    # other small loss and other teaches (soft label)

    loss_1_2 = my_cross_entropy(y_1, y_2.detach())
    loss_2_1 = my_cross_entropy(y_2, y_1.detach())
    loss_1_t = F.cross_entropy(y_1, t, reduction = 'none')
    loss_2_t = F.cross_entropy(y_2, t, reduction = 'none')
    _, ind_1_t_sorted = torch.sort(loss_1_t)
    _, ind_2_t_sorted = torch.sort(loss_2_t)
    
    remember_rate = 1 - forget_rate
    num_remember = int(remember_rate * len(t))

    # exchange
    loss_1_update = torch.cat((loss_1_t[ind_2_t_sorted[:num_remember]], loss_1_2[ind_2_t_sorted[num_remember:]]))
    loss_2_update = torch.cat((loss_2_t[ind_1_t_sorted[:num_remember]], loss_2_1[ind_1_t_sorted[num_remember:]]))

    return torch.mean(loss_1_update), torch.mean(loss_2_update)



def train(config, model1, model2, train_iter, dev_iter):
    start_time = time.time()
    model1.train()
    model2.train()
    optimizer1 = torch.optim.Adam(model1.parameters(), lr=config.learning_rate)
    optimizer2 = torch.optim.Adam(model2.parameters(), lr=config.learning_rate)
    # 学习率指数衰减，每次epoch：学习率 = gamma * 学习率
    # scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)
    # scheduler1 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer1, T_max=config.num_epochs)
    # scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer2, T_max=config.num_epochs)
    total_batch = 0  # 记录进行到多少batch
    dev_best_loss1 = float('inf')
    dev_best_loss2 = float('inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升
    writer = SummaryWriter(log_dir=config.log_path + '/' + time.strftime('%m-%d_%H.%M', time.localtime()))
    for epoch in range(config.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
        # scheduler1.step() # 学习率衰减
        # scheduler2.step() # 学习率衰减
        for i, (trains, labels) in enumerate(train_iter):
            
            outputs1 = model1(trains)
            outputs2 = model2(trains)
            model1.zero_grad()
            model2.zero_grad()
            # loss = F.cross_entropy(outputs, labels)
            loss1, loss2 = my_loss_coteaching(outputs1, outputs2, labels, 0.1)
            loss1.backward()
            loss2.backward()
            optimizer1.step()
            optimizer2.step()
            if total_batch % 100 == 0:
                improve = ''
                # 每多少轮输出在训练集和验证集上的效果
                true = labels.data.cpu()
                predic1 = torch.max(outputs1.data, 1)[1].cpu()
                train_acc1 = metrics.accuracy_score(true, predic1)

                predic2 = torch.max(outputs2.data, 1)[1].cpu()
                train_acc2 = metrics.accuracy_score(true, predic2)
                dev_acc1, dev_loss1 = evaluate(config, model1, dev_iter)
                dev_acc2, dev_loss2 = evaluate(config, model2, dev_iter)
                if dev_loss1 < dev_best_loss1:
                    dev_best_loss1 = dev_loss1
                    torch.save(model1.state_dict(), config.save_path)
                    improve = '*'
                    last_improve = total_batch
                if dev_loss2 < dev_best_loss2:
                    dev_best_loss2 = dev_loss2
                    torch.save(model2.state_dict(), config.save_path)
                    improve = '*'
                    last_improve = total_batch
                time_dif = get_time_dif(start_time)
                msg = 'Iter: {0:>6},  Train Loss1: {1:>5.5},  Train Acc1: {2:>6.5%}, Train Loss2: {3:>5.5},  Train Acc2: {4:>6.5%},\
                       Val Loss1: {5:>5.5},  Val Acc1: {6:>6.5%},  Val Loss2: {7:>5.5},  Val Acc2: {8:>6.5%},  Time: {9} {10}'
                print(msg.format(total_batch, loss1.item(), train_acc1,loss2.item(), train_acc2, dev_loss1, dev_acc1, dev_loss2, dev_acc2, time_dif, improve))
                #writer.add_scalar("loss/train", loss.item(), total_batch)
                #writer.add_scalar("loss/dev", dev_loss, total_batch)
                #writer.add_scalar("acc/train", 1, total_batch)
                #writer.add_scalar("acc/dev", dev_acc, total_batch)
                model1.train()
                model2.train()
            total_batch += 1
            '''
            if total_batch - last_improve > config.require_improvement:
                # 验证集loss超过1000batch没下降，结束训练
                print("No optimization for a long time, auto-stopping...")
                flag = True
                break
            '''
        if flag:
            break
    writer.close()
    test(config, model1, dev_iter)
    test(config, model2, dev_iter)


def test(config, model, test_iter):
    # test
    model.load_state_dict(torch.load(config.save_path))
    model.eval()
    start_time = time.time()
    test_acc, test_loss, test_report, test_confusion = evaluate(config, model, test_iter, test=True)
    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}'
    print(msg.format(test_loss, test_acc))
    print("Precision, Recall and F1-Score...")
    print(test_report)
    print("Confusion Matrix...")
    print(test_confusion)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)


def evaluate(config, model, data_iter, test=False):
    model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    with torch.no_grad():
        for texts, labels in data_iter:
            outputs = model(texts)
            loss = F.cross_entropy(outputs, labels)
            # loss = torch.mean(my_cross_entropy(outputs, labels))
            loss_total += loss
            labels = labels.data.cpu().numpy()
            predic = torch.max(outputs.data, 1)[1].cpu().numpy()
            labels_all = np.append(labels_all, labels)
            predict_all = np.append(predict_all, predic)

    acc = metrics.accuracy_score(labels_all, predict_all)
    if test:
        report = metrics.classification_report(labels_all, predict_all, target_names=config.class_list, digits=4)
        confusion = metrics.confusion_matrix(labels_all, predict_all)
        return acc, loss_total / len(data_iter), report, confusion
    return acc, loss_total / len(data_iter)
